{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dee6cbd4-3c38-4925-806c-daba1ffcda55",
   "metadata": {},
   "source": [
    "### 使用多GPU进行训练模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffd14365-a471-43a9-b711-eab190b5c262",
   "metadata": {},
   "source": [
    "- DDP的基本用法\n",
    "1. 使用 torch.distributed.init_process_group 初始化进程组 \n",
    "2. 使用 torch.nn.parallel.DistributedDataParallel 创建 分布式模型 \n",
    "3. 使用 torch.utils.data.distributed.DistributedSampler 创建 DataLoader \n",
    "4. 调整其他必要的地方(tensor放到指定device上，S/L checkpoint，指标计算等)\n",
    "6. 使用 torchrun 开始训练\n",
    "7. 514"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83cf4a77-edf1-4ac0-9fac-15bb85c64311",
   "metadata": {},
   "source": [
    "#### 初始化进程组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e934ae1-c6cc-4337-8c14-ef6316e8892f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-17T11:08:15.723152Z",
     "iopub.status.busy": "2023-12-17T11:08:15.722518Z",
     "iopub.status.idle": "2023-12-17T11:08:15.889775Z",
     "shell.execute_reply": "2023-12-17T11:08:15.888564Z",
     "shell.execute_reply.started": "2023-12-17T11:08:15.723091Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.distributed as dist\n",
    "import torch.nn as nn\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP\n",
    "import torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01d304f7-3942-40d9-8f2e-9c2748d4bdb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_distributed_mode(args):\n",
    "    # set up distributed device\n",
    "    args.rank = int(os.environ[\"RANK\"])\n",
    "    args.local_rank = int(os.environ[\"LOCAL_RANK\"])\n",
    "    torch.cuda.set_device(args.rank % torch.cuda.device_count())\n",
    "    dist.init_process_group(backend=\"nccl\")\n",
    "    args.device = torch.device(\"cuda\", args.local_rank)\n",
    "    print(args.device,'argsdevice')\n",
    "    args.NUM_gpu = torch.distributed.get_world_size()\n",
    "    print(f\"[init] == local rank: {args.local_rank}, global rank: {args.rank} ==\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0542554c-f239-4703-b294-aea3f4c885be",
   "metadata": {},
   "source": [
    "- 在train.py中修改对应部分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "448e924a-55cc-4d75-ab8d-e0ce07c80f02",
   "metadata": {},
   "outputs": [],
   "source": [
    "    # Initialize Multi GPU \n",
    "    if args.multi_gpu == True :\n",
    "        init_distributed_mode(args)\n",
    "    else: \n",
    "        # Use Single Gpu \n",
    "        os.environ['CUDA_VISIBLE_DEVICES'] = args.device_gpu\n",
    "        device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "        print(f'Using {device} device')\n",
    "        args.device = device   \n",
    "    #The learning rate is automatically scaled \n",
    "    # (in other words, multiplied by the number of GPUs and multiplied by the batch size divided by 32).\n",
    "    args.lr = args.lr * args.NUM_gpu * (args.batch_size / 32)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e9b9b72-3b2a-4815-bb5a-f008696a07ff",
   "metadata": {},
   "source": [
    "#### 创建分布式模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f09bba1c-fa64-40bc-bf5b-d23374f47b17",
   "metadata": {},
   "source": [
    "- 加载好model模型后创建分布式模型\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05d902f3-349a-4f38-ae35-9f27874d5894",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.cuda()\n",
    "if args.multi_gpu:\n",
    "    # DistributedDataParallel\n",
    "    ssd300 = DDP(model , device_ids=[args.local_rank], output_device=args.local_rank)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9192b663-b587-4e68-8d8c-660d88a69032",
   "metadata": {},
   "source": [
    "#### 创建Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20bfb871-85c0-46bc-87fb-abf5f58791da",
   "metadata": {},
   "outputs": [],
   "source": [
    "    train_dataset = COCODetection(root=args.data.DATASET_PATH,image_set='train2017', \n",
    "                        transform=SSDTransformer(dboxes))\n",
    "\n",
    "    val_dataset = COCODetection(root=args.data.DATASET_PATH,image_set='val2017', \n",
    "                        transform=SSDTransformer(dboxes, val=True))\n",
    "    \n",
    "    if args.multi_gpu:\n",
    "        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,shuffle=True)\n",
    "        val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)\n",
    "        train_shuffle = False\n",
    "    else:\n",
    "        train_sampler = None\n",
    "        val_sampler = None\n",
    "        train_shuffle = True\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset, args.batch_size,\n",
    "                                  num_workers=args.num_workers,\n",
    "                                  shuffle=train_shuffle, \n",
    "                                  sampler=train_sampler,\n",
    "                                  pin_memory=True)\n",
    "\n",
    "    val_loader = torch.utils.data.DataLoader(val_dataset,\n",
    "                                batch_size=args.batch_size,\n",
    "                                shuffle=False,  # Note: distributed sampler is shuffled :(\n",
    "                                sampler=val_sampler,\n",
    "                                num_workers=args.num_workers)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1eb52fc-b117-4c7e-9aa8-b02aafc4fd40",
   "metadata": {},
   "source": [
    "#### 一些注意事项"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08d72f14-663c-4cc1-81e9-7ceef9ec3d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.local_rank == 0:\n",
    "                log.logger.info(epoch, acc)\n",
    "# Save model\n",
    "if args.save and args.local_rank == 0:\n",
    "    print(\"saving model...\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "189c5e88-f0b5-460b-9786-9046dbbef9ce",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f124b66e-4d65-4f7f-9900-de7d055f6441",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db4d429f-70b6-49aa-b655-8e11487b85cb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
